import argparse
import torch

def zjm_no_bias():
    if opt.weight_in.endswith('.pt'): 
        weight_in = torch.load(opt.weight_in)['model']

    weight_out={}
    for key in weight_in.keys():
        if 'bias' in key:
            continue
        else:
            print(key)
            weight_out[key] = weight_in[key]

    torch.save({'model' : weight_out}, opt.weight_out)

    
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weight_in', type=str, default='/home/zhangjm/backup/yolov3_ReLU_fuse_voc/20200321.pt', help='path to input weights file')
    parser.add_argument('--weight_out', type=str, default='/home/zhangjm/backup/20210321/yolov3_ReLU_fuse_voc_no_bias.pt', help='output weights')

    opt = parser.parse_args()
    print(opt)

    zjm_no_bias()
